import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import os

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 第一个卷积层
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        # 第二个卷积层
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        # 池化层
        self.pool = nn.MaxPool2d(2, 2)
        # dropout层
        self.dropout = nn.Dropout2d(0.25)
        # 全连接层
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        # dropout层
        self.dropout2 = nn.Dropout(0.5)
        
    def forward(self, x):
        # 第一个卷积块
        x = self.pool(F.relu(self.conv1(x)))
        # 第二个卷积块
        x = self.pool(F.relu(self.conv2(x)))
        # dropout
        x = self.dropout(x)
        # 展平
        x = x.view(-1, 64 * 7 * 7)
        # 全连接层
        x = F.relu(self.fc1(x))
        x = self.dropout2(x)
        x = self.fc2(x)
        return x

def load_local_mnist(data_path="./mnist_data"):
    """
    从本地目录加载MNIST数据
    """
    print(f"尝试从本地目录 {data_path} 加载MNIST数据...")
    
    # 检查是否存在本地MNIST数据文件
    mnist_npz_path = os.path.join(data_path, "mnist.npz") if os.path.isdir(data_path) else data_path
    
    if os.path.exists(mnist_npz_path) and mnist_npz_path.endswith('.npz'):
        try:
            # 从本地NPZ文件加载数据
            with np.load(mnist_npz_path, allow_pickle=True) as f:
                x_train, y_train = f['x_train'], f['y_train']
                x_test, y_test = f['x_test'], f['y_test']
            print(f"成功从本地文件 {mnist_npz_path} 加载数据")
            
            # 转换为PyTorch张量
            x_train = torch.FloatTensor(x_train).unsqueeze(1) / 255.0  # 添加通道维度并归一化
            y_train = torch.LongTensor(y_train)
            x_test = torch.FloatTensor(x_test).unsqueeze(1) / 255.0
            y_test = torch.LongTensor(y_test)
            
            # 创建数据集
            train_dataset = TensorDataset(x_train, y_train)
            test_dataset = TensorDataset(x_test, y_test)
            
            return train_dataset, test_dataset
        except Exception as e:
            print(f"从本地NPZ文件加载数据失败: {e}")
    
    # 如果本地加载失败，使用torchvision方式加载
    print("使用torchvision方式加载MNIST数据...")
    import torchvision
    import torchvision.transforms as transforms
    
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    train_dataset = torchvision.datasets.MNIST(
        root='./data', train=True, download=True, transform=transform)
    test_dataset = torchvision.datasets.MNIST(
        root='./data', train=False, download=True, transform=transform)
    
    return train_dataset, test_dataset

def visualize_samples(train_loader, num_samples=9):
    """可视化数据样本"""
    # 获取一批数据
    dataiter = iter(train_loader)
    images, labels = next(dataiter)
    
    plt.figure(figsize=(8, 8))
    for i in range(min(num_samples, len(images))):
        plt.subplot(3, 3, i+1)
        # 反归一化图像以便显示
        img = images[i] / 2 + 0.5  # 反归一化
        npimg = img.numpy()
        plt.imshow(np.transpose(npimg, (1, 2, 0)).squeeze(), cmap='gray')
        plt.title(f'标签: {labels[i].item()}')
        plt.axis('off')
    plt.tight_layout()
    plt.show()

def train_model(model, train_loader, test_loader, epochs=5):
    """训练模型"""
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    # 记录训练历史
    train_losses = []
    train_accuracies = []
    test_accuracies = []
    
    # 训练模型
    for epoch in range(epochs):
        model.train()  # 设置为训练模式
        running_loss = 0.0
        correct_train = 0
        total_train = 0
        
        for i, (images, labels) in enumerate(train_loader):
            # 将数据移到设备上
            images, labels = images.to(device), labels.to(device)
            
            # 清零梯度
            optimizer.zero_grad()
            
            # 前向传播
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            # 反向传播
            loss.backward()
            
            # 更新参数
            optimizer.step()
            
            # 统计信息
            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total_train += labels.size(0)
            correct_train += (predicted == labels).sum().item()
            
            # 每100个批次打印一次信息
            if (i + 1) % 100 == 0:
                print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')
        
        # 计算训练准确率
        train_acc = 100 * correct_train / total_train
        avg_loss = running_loss / len(train_loader)
        
        # 计算测试准确率
        test_acc = evaluate_model(model, test_loader)
        
        # 记录历史
        train_losses.append(avg_loss)
        train_accuracies.append(train_acc)
        test_accuracies.append(test_acc)
        
        print(f'Epoch [{epoch+1}/{epochs}], 平均损失: {avg_loss:.4f}, 训练准确率: {train_acc:.2f}%, 测试准确率: {test_acc:.2f}%')
    
    return train_losses, train_accuracies, test_accuracies

def evaluate_model(model, test_loader):
    """评估模型"""
    model.eval()  # 设置为评估模式
    correct = 0
    total = 0
    
    with torch.no_grad():  # 不计算梯度
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    accuracy = 100 * correct / total
    return accuracy

def plot_training_history(train_losses, train_accuracies, test_accuracies):
    """绘制训练历史"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    # 绘制损失曲线
    ax1.plot(train_losses, label='训练损失')
    ax1.set_title('模型损失')
    ax1.set_xlabel('轮次')
    ax1.set_ylabel('损失')
    ax1.legend()
    ax1.grid(True)
    
    # 绘制准确率曲线
    ax2.plot(train_accuracies, label='训练准确率')
    ax2.plot(test_accuracies, label='测试准确率')
    ax2.set_title('模型准确率')
    ax2.set_xlabel('轮次')
    ax2.set_ylabel('准确率 (%)')
    ax2.legend()
    ax2.grid(True)
    
    plt.tight_layout()
    plt.show()

def test_predictions(model, test_loader, num_samples=10):
    """测试预测结果"""
    model.eval()
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    
    # 将数据移到设备上
    images, labels = images.to(device), labels.to(device)
    
    # 预测
    outputs = model(images[:num_samples])
    _, predicted = torch.max(outputs, 1)
    
    # 显示结果
    print("预测结果示例:")
    print("样本\t真实标签\t预测标签")
    print("-" * 30)
    for i in range(num_samples):
        print(f"{i+1}\t{labels[i].item()}\t\t{predicted[i].item()}")
    
    # 可视化预测结果
    plt.figure(figsize=(12, 6))
    for i in range(min(10, num_samples)):
        plt.subplot(2, 5, i+1)
        # 反归一化图像以便显示
        img = images[i].cpu() / 2 + 0.5
        npimg = img.numpy()
        plt.imshow(np.transpose(npimg, (1, 2, 0)).squeeze(), cmap='gray')
        color = 'green' if predicted[i] == labels[i] else 'red'
        plt.title(f'真实: {labels[i].item()}, 预测: {predicted[i].item()}', color=color)
        plt.axis('off')
    
    plt.tight_layout()
    plt.savefig('pytorch_mnist_predictions.png', dpi=300, bbox_inches='tight')
    print("\n预测结果已保存为 'pytorch_mnist_predictions.png'")
    plt.show()

def main():
    """主函数"""
    print("MNIST手写数字识别 - PyTorch CNN实现")
    print("=" * 50)
    
    # 加载数据
    print("\n加载MNIST数据...")
    train_dataset, test_dataset = load_local_mnist("../mnist_data")
    
    # 创建数据加载器
    batch_size = 64
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    
    # 可视化样本数据
    print("\n可视化样本数据...")
    visualize_samples(train_loader)
    
    # 创建模型
    print("\n创建CNN模型...")
    model = CNN().to(device)
    
    # 显示模型结构
    print(model)
    
    # 训练模型
    print("\n开始训练模型...")
    train_losses, train_accuracies, test_accuracies = train_model(model, train_loader, test_loader, epochs=5)
    
    # 绘制训练历史
    print("\n绘制训练历史...")
    plot_training_history(train_losses, train_accuracies, test_accuracies)
    
    # 测试预测
    print("\n测试预测结果...")
    test_predictions(model, test_loader)
    
    # 保存模型
    torch.save(model.state_dict(), '../mnist_cnn_model.pth')
    print("\n模型已保存为 'mnist_cnn_model.pth'")
    
    # 最终评估
    final_accuracy = evaluate_model(model, test_loader)
    print(f"\n最终测试准确率: {final_accuracy:.2f}%")

if __name__ == "__main__":
    main()